import warnings

import pandas as pd
import statsmodels.api as sm


def predict(file_path, sheet_name, which_col):
    # plt.rcParams['figure.figsize'] = (20.0, 10.0)
    # plt.rcParams.update({'font.size': 12})
    # plt.style.use('ggplot')

    # Load the data
    data = pd.read_excel(file_path, sheet_name=sheet_name, index_col='时刻')
    data.index = pd.to_datetime(data.index, unit='h', origin=pd.Timestamp('2019-01-01'))
    data = data[which_col]

    # Plot the data
    # data.plot()

    # 可根据日期索引取用不同的数据作为训练集和测试集
    # train_data = data['2019-01-01':'2019-01-31']

    warnings.filterwarnings("ignore")  # specify to ignore warning messages

    # 将训练数据送入best model中
    mod = sm.tsa.statespace.SARIMAX(data,
                                    order=[3, 1, 3],
                                    seasonal_order=[2, 1, 1, 24],
                                    enforce_stationarity=False,
                                    enforce_invertibility=False)

    # 拟合
    results = mod.fit()
    pred2 = results.get_forecast(steps=24)
    result_list = [pred2.predicted_mean[i] for i in range(24)]
    return result_list
